# Copyright 2021-2023 @ Shenzhen Bay Laboratory &
#                       Peking University &
#                       Huawei Technologies Co., Ltd
#
# This code is a part of MindSPONGE:
# MindSpore Simulation Package tOwards Next Generation molecular modelling.
#
# MindSPONGE is open-source software based on the AI-framework:
# MindSpore (https://www.mindspore.cn/)
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""
Callback to write H5MD trajectory file
"""

import signal
import sys
from typing import List
import numpy as np

from ..system import Molecule
from ..optimizer import Updater
from ..data.export import H5MD
from ..system.modelling.pdb_generator import gen_pdb


class WriteH5MD:

    def __init__(self,
                 system: Molecule,
                 filename: str,
                 save_freq: int = 1,
                 constant_volume: bool = True,
                 directory: str = None,
                 mode: str = 'w',
                 write_velocity: bool = False,
                 write_force: bool = False,
                 write_image: bool = True,
                 write_metrics: bool = None,
                 write_bias: bool = False,
                 length_unit: str = None,
                 energy_unit: str = None,
                 dtype: str = 'float32',
                 compression: str = 'gzip',
                 compression_opts: int = 4,
                 auto_close: bool = True,
                 save_last_pdb: str = None
                 ):

        if mode not in ['w', 'w-', 'x', 'a']:
            raise ValueError(f'WriteH5MD only support "w", "w-", "x" and "a" but got: {mode}')

        self.system = system
        self.num_walker = system.num_walker
        self.units = system.units
        self.filename = filename
        self.h5md = H5MD(self.system, filename, directory, mode=mode,
                         length_unit=length_unit, energy_unit=energy_unit,
                         compression=compression, compression_opts=compression_opts)

        self.convert_to_angstram = self.system.units.convert_length_to('A')
        self.use_pbc = True
        self.const_volume = constant_volume
        self.dtype = dtype
        self.save_freq = save_freq
        self.auto_close = auto_close

        self.potential = 0
        self.energies = 0
        self.kinetics = 0
        self.tot_energy = 0
        self.temperature = 0
        self.pressure = 0
        self.volume = 0
        self.bias = 0
        self.biases = 0

        self.write_image = system.pbc_box is not None
        self.write_velocity = write_velocity
        self.write_force = write_force

        if save_last_pdb is None:
            self.last_pdb_name = None
            self.save_pdb = False
        else:
            self.save_pdb = True
            self.last_pdb_name = save_last_pdb

        if mode == 'a':
            self.init_h5md = False
        else:
            self.init_h5md = True
            if self.use_pbc and self.write_image:
                self.h5md.set_image()

            if self.write_velocity:
                self.h5md.set_velocity()

            if self.write_force:
                self.h5md.set_force()

            self.h5md.add_observables('potential_energy', (), self.dtype, self.units.energy_unit_name)
            self.h5md.add_observables('total_energy', (), self.dtype, self.units.energy_unit_name)

        self.use_updater = True

        self.write_metrics = write_metrics
        self.write_energies = False
        self.write_bias = write_bias

        self.count = 0
        self.count_records = 0

        # Detect process kill signal and save the data.
        signal.signal(signal.SIGINT, self.__stop__)
        signal.signal(signal.SIGTERM, self.__stop__)

    def __enter__(self):
        """Return the enter target."""
        return self

    def __exit__(self, *err):
        """Release resources here if have any."""

    def __stop__(self, signal_, frame_):
        """
        Save data when process killed.
        """
        # pylint: disable=unused-argument
        print(f'\n\033[33mProgram process terminated. {self.count_records} steps saved in H5MD file.\033[0m\n')
        self.close()
        sys.exit(0)

    def close(self):
        if self.save_pdb:
            self.save_to_pdb()
        self.h5md.close()
        return self

    def begin(self, num_energies: int, energy_names: List[str], num_biases: int, bias_names: List[str], metrics, grids, hills):
        """
        Called once before the network executing.

        Args:
            run_context (RunContext): Include some information of the model.
        """

        def _init_energies(num_energies: int, energy_names: List[str]):
            if num_energies == 0:
                return

            self.write_energies = True
            if self.init_h5md:
                energy_names = [name.encode('ascii', 'ignore') for name in energy_names]
                energies = self.h5md.add_observables(
                    'energies', (num_energies,), self.dtype, self.units.energy_unit_name)
                self.h5md.create_dataset(energies, 'labels', shape=(1, num_energies), data=energy_names)
            return

        def _init_biases(num_biases: int, bias_names: List[str], grids, hills):
            if num_biases == 0:
                return

            self.write_bias = True
            if self.init_h5md:
                self.h5md.add_observables('bias_potential', (), self.dtype, self.units.energy_unit_name)
                bias_names = [name.encode('ascii', 'ignore') for name in bias_names]
                biases = self.h5md.add_observables(
                    'biases', (num_biases,), self.dtype, self.units.energy_unit_name)
                self.h5md.create_dataset(biases, 'labels', shape=(1, num_biases), data=bias_names)
                meta_potential = self.h5md.parameters.create_group('meta_potential')
                if hills is not None:
                    for i in range(len(hills)):
                        meta_potential.create_dataset(f'hills{i}',data=hills[i])
                        meta_potential.create_dataset(f'grids{i}',data=grids[i])
            return

        _init_energies(num_energies, energy_names)
        _init_biases(num_biases, bias_names, grids, hills)

        if self.use_updater:
            self.h5md.add_observables('kinetic_energy', (), self.dtype, self.units.energy_unit_name)
            self.h5md.add_observables('temperature', (), self.dtype, 'K')
            if self.init_h5md:
                self.h5md.add_observables('pressure', (), self.dtype, 'bar')
                self.h5md.add_observables('volume', (), self.dtype, self.units.volume_unit_name)
                self.h5md.set_box(self.const_volume)
        # else:
        #     self.use_updater = False
        #     if self.use_pbc and self.init_h5md:
        #         self.h5md.set_box(True)
        #     if self.write_velocity and not isinstance(cb_params.optimizer, Updater):
        #         self.write_velocity = False
        #         print(f'Warning! The optimizer "{cb_params.optimizer}" does not has the attribute "velocity".')

        # if cb_params.metrics is None:
        #     if self.write_metrics:
        #         print('[WARNING] No Metrics found. Metrics will not be output to H5MD file.')
        #     self.write_metrics = False
        # else:
        #     if self.write_metrics is None:
        #         self.write_metrics = True
        #     if self.init_h5md:
        #         for name, shape, unit in \
        #             zip(cb_params.metrics, cb_params.metrics_shape, cb_params.metrics_units):
        #             shape_ = () if shape == (1,) else shape
        #             self.h5md.add_observables(name, shape_, self.dtype, unit)

        if self.write_metrics:
            if self.init_h5md:
                for name, cv in metrics.items():
                    self.h5md.add_observables(name, cv.shape, self.dtype, cv._unit)

        self.init_h5md = False

    def write_info(self, step, time, coordinate, velocity, force, potential, energies, kinetics, temperature, pbc_box, pressure, bias, biases, hills, metrics):
        """
        Called before each step beginning.

        Args:
            run_context (RunContext): Include some information of the model.
        """

        self.h5md.write_step(step)
        self.h5md.write_time(time.detach().cpu().numpy())

        self.potential = potential.detach().cpu().numpy().squeeze()
        self.h5md.write_observables('potential_energy', self.potential)

        if self.write_energies:
            self.energies = energies.detach().cpu().numpy().squeeze()
            self.h5md.write_observables('energies', self.energies)

        if self.write_bias:
            self.bias = bias.detach().cpu().numpy().squeeze()
            self.h5md.write_observables('bias_potential', self.bias)
            self.biases = biases.detach().cpu().numpy().squeeze()
            self.h5md.write_observables('biases', self.biases)
            if hills is not None:
                for i in range(len(hills)):
                    self.h5md.hdf5_file['parameters']['meta_potential'][f'hills{i}'][:] = hills[i]

        self.tot_energy = self.potential

        if self.use_updater:
            self.kinetics = kinetics.detach().cpu().numpy().squeeze()
            self.h5md.write_observables('kinetic_energy', self.kinetics)
            self.temperature = temperature.detach().cpu().numpy().squeeze()
            self.h5md.write_observables('temperature', self.temperature)
            self.tot_energy += + self.kinetics
        self.h5md.write_observables('total_energy', self.tot_energy)    
        coordinate = coordinate.detach().cpu().numpy().squeeze()
        self.h5md.write_position(coordinate)

        if self.use_pbc:
            if not self.const_volume:
                self.h5md.write_box(pbc_box.detach().cpu().numpy().squeeze())
            if self.write_image:
                image = self.system.calc_image().detach().cpu().numpy().squeeze()
                self.h5md.write_image(image)
            if pressure is not None:
                self.pressure = pressure.detach().cpu().numpy().squeeze()
                self.h5md.write_observables('pressure', self.pressure)
                self.volume = np.prod(pbc_box.detach().cpu().numpy().squeeze(),-1)
                self.h5md.write_observables('volume', self.volume)

        if self.write_metrics:
            for metric, value in metrics.items():
                self.h5md.write_observables(metric, value.detach().cpu().numpy().squeeze())

        if self.write_velocity:
            velocity = velocity.detach().cpu().numpy().squeeze()
            self.h5md.write_velocity(velocity)
        if self.write_force:
            force = force.detach().cpu().numpy().squeeze()
            self.h5md.write_force(force)

        self.h5md.hdf5_file.flush()


    def end(self, hills):
        """
        Called once after network training.

        Args:
            run_context (RunContext): Include some information of the model.
        """
        #pylint: disable=unused-argument
        if hills is not None:
            for i in range(len(hills)):
                self.h5md.hdf5_file['parameters']['meta_potential'][f'hills{i}'][:] = hills[i]

        if self.auto_close:
            self.close()

    def save_to_pdb(self):
        """ Save the system information into a pdb file.
        """
        last_resname = self.system.residue_name
        if len(last_resname[0]) == 4:
            last_resname[0] = last_resname[0][1:]
        if len(last_resname[-1]) == 4:
            last_resname[-1] = last_resname[-1][1:]
        gen_pdb(self.system.coordinate.detach().cpu().numpy() * self.convert_to_angstram,
                self.system.atom_name[0],
                np.take(last_resname, self.system.atom_resid),
                self.system.atom_resid.numpy() + 1,
                pdb_name=self.last_pdb_name)
